package cn.sciento.transfer.infra.datasource.utils;

import cn.sciento.core.exception.CommonException;import java.sql.Connection;
import java.sql.SQLException;
import cn.sciento.transfer.infra.datasource.factory.DataSourceContextHolder;
import cn.sciento.transfer.infra.datasource.provider.TransferDataSourceProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.TransactionTemplate;
import org.springframework.util.StringUtils;

@Component
public class TransferDataSourceTemplate {
  private static final Logger LOGGER = LoggerFactory.getLogger(TransferDataSourceTemplate.class);
  
  @Autowired
  private JdbcTemplate jdbcTemplate;
  
  @Autowired
  private TransactionTemplate transactionTemplate;
  
  public <T> T doExecute(String dataSourceCode, String serviceName, JdbcExecution<T> jdbcExecution) {
    TransferDataSourceProvider.loadDataSource(dataSourceCode, serviceName);
    DataSourceContextHolder.setDataSourceLookupKey(TransferDataSourceProvider.getDataSourceKey(dataSourceCode));
    try {
      return jdbcExecution.doExecuteInternal(this.jdbcTemplate);
    } catch (Exception ex) {
      LOGGER.debug("execute failed: ", ex);
      DataSourceContextHolder.clearDataSourceLookupKey();
      throw new CommonException(ex);
    } finally {
      DataSourceContextHolder.clearDataSourceLookupKey();
    } 
  }
  
  public <T> T doExecuteInConnection(String dataSourceCode, String serviceName, final String dbCatalog, final JdbcConnectionExecution<T> jdbcConnectionExecution) {
    TransferDataSourceProvider.loadDataSource(dataSourceCode, serviceName);
    DataSourceContextHolder.setDataSourceLookupKey(TransferDataSourceProvider.getDataSourceKey(dataSourceCode));
    try {
      T result = (T)this.jdbcTemplate.execute((ConnectionCallback<T>) con -> {
        if (!StringUtils.isEmpty(dbCatalog))
          con.setCatalog(dbCatalog);
        return jdbcConnectionExecution.doExecuteInConnection(con);
      });
      return result;
    } catch (Exception ex) {
      LOGGER.debug("execute failed: ", ex);
      DataSourceContextHolder.clearDataSourceLookupKey();
      throw new CommonException(ex);
    } finally {
      DataSourceContextHolder.clearDataSourceLookupKey();
    } 
  }
  
  public <T> T doExecuteInTransaction(String dataSourceCode, String serviceName, JdbcExecution<T> jdbcExecution) {
    TransferDataSourceProvider.loadDataSource(dataSourceCode, serviceName);
    DataSourceContextHolder.setDataSourceLookupKey(TransferDataSourceProvider.getDataSourceKey(dataSourceCode));
    try {
      return (T)this.transactionTemplate.execute(status -> jdbcExecution.doExecuteInternal(this.jdbcTemplate));
    } catch (Exception ex) {
      LOGGER.debug("execute failed: ", ex);
      DataSourceContextHolder.clearDataSourceLookupKey();
      throw new CommonException(ex);
    } finally {
      DataSourceContextHolder.clearDataSourceLookupKey();
    } 
  }
  
  public <T> T doExecuteInTransactionWithConnection(String dataSourceCode, String serviceName, final String dbCatalog, final JdbcConnectionExecution<T> jdbcConnectionExecution) {
    TransferDataSourceProvider.loadDataSource(dataSourceCode, serviceName);
    DataSourceContextHolder.setDataSourceLookupKey(TransferDataSourceProvider.getDataSourceKey(dataSourceCode));
    try {
      return (T)this.transactionTemplate.execute(status -> this.jdbcTemplate.execute(new ConnectionCallback<T>() {
              public T doInConnection(Connection con) throws SQLException, DataAccessException {
                if (!StringUtils.isEmpty(dbCatalog))
                  con.setCatalog(dbCatalog); 
                return jdbcConnectionExecution.doExecuteInConnection(con);
              }
            }));
    } catch (Exception ex) {
      LOGGER.debug("execute failed: ", ex);
      DataSourceContextHolder.clearDataSourceLookupKey();
      throw new CommonException(ex);
    } finally {
      DataSourceContextHolder.clearDataSourceLookupKey();
    } 
  }
  
  @FunctionalInterface
  public static interface JdbcExecution<T> {
    T doExecuteInternal(JdbcTemplate param1JdbcTemplate);
  }
  
  @FunctionalInterface
  public static interface JdbcConnectionExecution<T> {
    T doExecuteInConnection(Connection param1Connection);
  }
}
